#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under theglim Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

from scipy.special import dawsn
import scipy.ndimage

from numpy import dot

# from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors
from camori_harris import PictureHarris

from quaternion import Quat, random_quaternion

import simplejson

import scipy.optimize
#from fmin_sa import Sampling
from simplex_sa import SimplexSO3

import Image

from plot_aux import *

set_printoptions(precision=7)



zernike_V11_11 = array([[0., 0., -0.0022800053164889573, -0.007185706518456357, -0.005318695380781893, 0., 0.005318695380781894, 0.007185706518456347, 0.00228000531648896, 0., 0.],
[0., -0.006621652628393973, -0.017169982733489353, -0.012021036814425248, -0.006010518407212622, 0., 0.006010518407212628, 0.012021036814425255, 0.01716998273348936, 0.006621652628393962, 0.],
[-0.003877338545822786, -0.02291955601966685, -0.018031555221637872, -0.012021036814425248, -0.006010518407212622, 0., 0.006010518407212628, 0.012021036814425255, 0.01803155522163788, 0.022919556019666845, 0.0038773385458227825],
[-0.017781116954670682, -0.0240420736288505, -0.018031555221637872, -0.012021036814425248, -0.006010518407212622, 0., 0.006010518407212628, 0.012021036814425255, 0.01803155522163788, 0.024042073628850504, 0.017781116954670682],
[-0.026796894565489632, -0.024042073628850497, -0.018031555221637872, -0.012021036814425248, -0.006010518407212622, 0., 0.006010518407212628, 0.012021036814425252, 0.018031555221637876, 0.024042073628850497, 0.026796894565489615],
[-0.029802153769095924, -0.024042073628850497, -0.018031555221637872, -0.012021036814425248, -0.0060105184072126215, 0., 0.006010518407212627, 0.012021036814425252, 0.018031555221637876, 0.024042073628850497, 0.02980215376909591],
[-0.026796894565489626, -0.024042073628850497, -0.018031555221637872, -0.012021036814425248, -0.006010518407212622, 0., 0.006010518407212628, 0.012021036814425252, 0.018031555221637876, 0.024042073628850497, 0.026796894565489598],
[-0.017781116954670686, -0.0240420736288505, -0.018031555221637872, -0.012021036814425248, -0.006010518407212622, 0., 0.006010518407212628, 0.012021036814425255, 0.01803155522163788, 0.024042073628850504, 0.017781116954670665],
[-0.0038773385458227873, -0.02291955601966685, -0.018031555221637872, -0.012021036814425248, -0.006010518407212622, 0., 0.006010518407212628, 0.012021036814425255, 0.01803155522163788, 0.022919556019666845, 0.0038773385458227825],
[0., -0.0066216526283939675, -0.017169982733489353, -0.012021036814425248, -0.006010518407212622, 0., 0.006010518407212628, 0.012021036814425255, 0.017169982733489353, 0.0066216526283939545, 0.],
[0., 0., -0.0022800053164889577, -0.007185706518456362, -0.005318695380781892, 0., 0.005318695380781895, 0.007185706518456365, 0.002280005316488954, 0., 0.]])






def all_oris(patch, a_i, sigma_i):
    Nfe = a_i.shape[0]
    Nf = 360*2
    deltatheta = 2*pi/(Nf)

    qqp = 1j*zeros(Nf)
    qqn = 1j*zeros(Nf)

    out = 1j*zeros(Nf)
    for m in range(Nfe):
        qqp[m] = (sigma_i[m] * a_i[m] * patch).sum()
        qqn[m] = ((-1)**m * sigma_i[m] * conj(a_i[m]) * patch).sum()

    for n in range(Nf):
        theta = n*deltatheta

        # qq = 1j*zeros((Nw,Nw))
        # for m in range(Nfe):
        #     qq += sigma_i[m] * a_i[m] * exp(1j * -m * theta)
        #     qq += (-1)**m * sigma_i[m] * conj(a_i[m]) * exp(1j * m * theta)
        # out[n] = (qq * patch).sum()

        for m in range(Nfe):
            out[n] += qqp[m] * exp(1j * -m * theta)
            out[n] += qqn[m] * exp(1j * m * theta)

    return out

def ori_target(theta,qqp,qqn):
    out = 1j * 0
    for m in range(Nfe):
        out += qqp[m] * exp(1j * -m * theta)
        out += qqn[m] * exp(1j * m * theta)
    return -(real(out)**2+imag(out)**2)
    
def find_ori(patch, a_i, sigma_i, tini=0.0):
    Nfe = a_i.shape[0]

    qqp = 1j*zeros(Nfe)
    qqn = 1j*zeros(Nfe)

    for m in range(Nfe):
        qqp[m] = (sigma_i[m] * a_i[m] * patch).sum()
        qqn[m] = ((-1)**m * sigma_i[m] * conj(a_i[m]) * patch).sum()

    ang = scipy.optimize.brent(ori_target, args=(qqp,qqn), brack=(tini-.1,tini+.1))
    
    return -ang



def myf(x):
    return 4.0/sqrt(pi) * (-x -dawsn(x) +2*x**2*dawsn(x))



if __name__ == '__main__' :

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load image and initialize pic object

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])
    filename = fileroot+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    im = imr.sum(2)/3

    Nfe = 12

    Nw = 11
    sigma1 = 3.0
    sigma2 = 3.0
    yy,xx = mgrid[0.:Nw,0.:Nw] / floor(Nw/2) - 1.0
    myfilt = (exp(-(sigma2*yy)**2) * exp(-(sigma1*xx)**2) * (4 * (sigma1*xx)**2 - 2)) \
                + 1j * (exp(-(sigma2*yy)**2) * myf(sigma1*xx))

    edge_x = scipy.ndimage.convolve(im/256., real(myfilt)) + 1j * scipy.ndimage.convolve(im/256., imag(myfilt))
    edge_y = scipy.ndimage.convolve(im/256., real(myfilt).T) + 1j * scipy.ndimage.convolve(im/256., imag(myfilt).T)
    Ix = real(edge_x)**2 + imag(edge_x)**2
    Iy = real(edge_y)**2 + imag(edge_y)**2


    edgels = []

    gstep = 16

    for j in range(1,Ix.shape[0]-1, gstep):
        for k in range(1,Ix.shape[1]-1):
            #if (Ix[j,k]+Iy[j,k]) < 3.0:
            if (Ix[j,k]+Iy[j,k]) < 0.25:
                continue
            if Ix[j,k] >= Iy[j,k]:
                if Ix[j,k] > Ix[j,k+1] and Ix[j,k] >= Ix[j,k-1]:
                    edgels.append(array([k,j,0,1.0]))

    for j in range(1,Ix.shape[0]-1):
        for k in range(1,Ix.shape[1]-1,gstep):
            #if (Ix[j,k]+Iy[j,k]) < 3.0:
            if (Ix[j,k]+Iy[j,k]) < 0.5:
                continue
            if Ix[j,k] < Iy[j,k]:
                if Iy[j,k] > Iy[j+1,k] and Iy[j,k] >= Iy[j-1,k]:
                    edgels.append(array([k,j,1.0,0]))


    read_perona = np.load('perona_bases.npz')
    a_i = read_perona['a_i']
    sigma_i = read_perona['sigma_i']
    #a_i[0]-=mean(a_i[0])

    edgels = array(edgels)

    Ne = len(edgels)
    remove = []
    for n in range(Ne):
        ex,ey,l1,l2 = edgels[n]
        if (ex < (Nw/2+1) or ex > (im.shape[1]-Nw/2-1) or
            ey < (Nw/2+1) or ey > (im.shape[0]-Nw/2-1)):
            remove.append(n)
            continue
        patch = im[ey-Nw/2:ey+Nw/2+1,ex-Nw/2:ex+Nw/2+1]
        tt = find_ori(patch, a_i, sigma_i)
        edgels[n,2] = cos(tt)
        edgels[n,3] = sin(tt)
    edgels = edgels[ array([not x in remove for x in range(Ne)]) ]
    Ne = edgels.shape[0]

    # figure(1)
    # for n in range(8):
    #     subplot(4,4,n+1)
    #     imshow(real(a_i[n]), vmin=-.01, vmax=.01)
    #     subplot(4,4,n+9)
    #     imshow(imag(a_i[n]), vmin=-.01, vmax=.01)

    figure(2)
    imshow(im, cmap=cm.bone)
    vv = 10*array([-1,1])
    for n in range(Ne):
        px,py,ux,uy = edgels[n]
        plot(px,py, 'r+')
        #plot(px+ux*vv, py+uy*vv, 'r-', lw=2)
        plot(px+uy*vv, py-ux*vv, 'r-', lw=2)



    figure(3)
    subplot(2,2,1)
    imshow(real(edge_x), vmin=-2, vmax=2)
    subplot(2,2,2)
    imshow(imag(edge_x), vmin=-2, vmax=2)
    subplot(2,2,3)
    imshow(real(edge_y), vmin=-2, vmax=2)
    subplot(2,2,4)
    imshow(imag(edge_y), vmin=-2, vmax=2)


    figure(4)
    subplot(2,2,1)
    imshow(Ix, vmin=0, vmax=10, cmap=cm.bone)
    subplot(2,2,2)
    imshow(Iy, vmin=0, vmax=10, cmap=cm.bone)
    subplot(2,2,3)
    imshow(Ix+Iy, vmin=0, vmax=10, cmap=cm.bone)
    
    # figure(6)
    # subplot(1,2,1)
    # imshow(real(myfilt), vmin=-2, vmax=2)
    # subplot(1,2,2)
    # imshow(imag(myfilt), vmin=-2, vmax=2)

